{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 361,
   "id": "c7585bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 362,
   "id": "997f4969",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "eps = 0.0001\n",
    "init = -1\n",
    "\n",
    "A = np.zeros((n, n))\n",
    "A[np.arange(n - 1), np.arange(n - 1) + 1] = 1\n",
    "A[np.arange(n - 1) + 1, np.arange(n - 1)] = 2\n",
    "A[np.arange(n), np.arange(n)] = 3\n",
    "\n",
    "B = np.zeros((n, n))\n",
    "B[np.arange(n - 2), np.arange(n - 2) + 2] = 1\n",
    "B[np.arange(n - 2) + 2, np.arange(n - 2)] = 3\n",
    "B[np.arange(n), np.arange(n)] = 4\n",
    "B[:, 2] = 0\n",
    "B[2, :] = 0\n",
    "\n",
    "m = 4\n",
    "\n",
    "C = np.zeros((m, m))\n",
    "C[np.arange(m - 1), np.arange(m - 1) + 1] = 1\n",
    "C[np.arange(m - 1) + 1, np.arange(m - 1)] = 1\n",
    "C[np.arange(m), np.arange(m)] = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 363,
   "id": "dfbe5680",
   "metadata": {},
   "outputs": [],
   "source": [
    "def abort(err):\n",
    "    errdict = {\n",
    "        1 : \"The sizes of matrixs are not matched!\",\n",
    "        2 : \"The size of matrix is wrong!\"\n",
    "    }\n",
    "    return errdict.get(err,\"error!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 364,
   "id": "fc2c4547",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Compress(A):\n",
    "    n = np.size(A, 0)\n",
    "    m = np.size(A, 1)\n",
    "    k = 0\n",
    "    row = [k]\n",
    "    col = []\n",
    "    val = []\n",
    "    for i in np.arange(n):\n",
    "        for j in np.arange(m):\n",
    "            if A[i, j] > eps or A[i, j] < -eps:\n",
    "                k = k + 1\n",
    "                col.append(j)\n",
    "                val.append(A[i, j])\n",
    "        row.append(k)\n",
    "    return row, col, val\n",
    "#无零行or零列时，按此压缩"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 366,
   "id": "48a73ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Compress_initial(A):\n",
    "    n = np.size(A, 0)\n",
    "    m = np.size(A, 1)\n",
    "    if n >= m :\n",
    "        for i in np.arange(m):\n",
    "            if(A[i, i] < eps):\n",
    "                A[i, i] = init\n",
    "            \n",
    "        for i in np.arange(m, n):\n",
    "            if(A[i, m] < eps):\n",
    "                A[i, m] = init\n",
    "    else:\n",
    "        for i in np.arange(n):\n",
    "            if(A[i, i] < eps):\n",
    "                A[i, i] = init\n",
    "            \n",
    "        for i in np.arange(n, m):\n",
    "            if(A[n, i] < eps):\n",
    "                A[n, i] = init\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 367,
   "id": "388e6b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Compress_finish(A_CSR):\n",
    "    l = [i for i, x in enumerate(A_CSR[2]) if x == init]\n",
    "    if -1 in l:\n",
    "        return A_CSR\n",
    "    \n",
    "    for i in l:\n",
    "        A_CSR[2][i] = 0.0\n",
    "        \n",
    "    return A_CSR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 365,
   "id": "be900240",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Compress_universal(A):\n",
    "    A = Compress_initial(A)\n",
    "    A_CSR = Compress(A)\n",
    "    A_CSR = Compress_finish(A_CSR)\n",
    "    return A_CSR\n",
    "#存在零行or零列时，按此压缩，当然无零行or零列也可按此压缩，但效率低了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 372,
   "id": "5611b9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unCompress(A_CSR, n, m):\n",
    "    a = np.size(A_CSR[0]) - 1\n",
    "    b = np.max(A_CSR[1]) + 1\n",
    "    \n",
    "    if n <= 0 or m <= 0 or a > n or b > m:\n",
    "        return abort(2)\n",
    "    \n",
    "    A = np.zeros((n, m))\n",
    "    for i in np.arange(a):\n",
    "        for j in np.arange(A_CSR[0][i + 1] - A_CSR[0][i]):\n",
    "            A[i][A_CSR[1][A_CSR[0][i] + j]] = A_CSR[2][A_CSR[0][i] + j]\n",
    "            \n",
    "    return A\n",
    "#由于可能有零行or零列，所以需要给出行数和列数，如果小于按行存储指标，则err，大于则自动添零行和零列\n",
    "#由compress_universal压缩的矩阵一定不会有零行和零列的情况，因为对角线默认存入CSR中，即使它的元素为0，然而由计算得到的矩阵可能存在最后一行\n",
    "#or列均为0，因此需要加上参数n和m确定大小，事实上将这些打包合并后不需要n与m，因为大小可以由原矩阵得出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 368,
   "id": "d7b61611",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Transfer1(A):\n",
    "    A_CSR = Compress(A)\n",
    "    n = np.size(A_CSR[0]) - 1\n",
    "    m = np.max(A_CSR[1]) + 1\n",
    "    T = np.zeros((m, n))\n",
    "    \n",
    "    for i in np.arange(n):\n",
    "        for j in np.arange(A_CSR[0][i + 1] - A_CSR[0][i]):\n",
    "            T[A_CSR[1][A_CSR[0][i] + j]][i] = A_CSR[2][A_CSR[0][i] + j]\n",
    "            \n",
    "    return Compress(T)\n",
    "#版本1，效率低"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 369,
   "id": "79834b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Transfer2(A):\n",
    "    A_CSR = Compress(A)\n",
    "    n = np.size(A_CSR[0]) - 1\n",
    "    m = np.max(A_CSR[1]) + 1\n",
    "    a = []\n",
    "    t = 0\n",
    "    trow = [t]\n",
    "    tcol = [[]for i in range(n)]\n",
    "    tval = [[]for i in range(n)]\n",
    "    \n",
    "    for i in np.arange(m):\n",
    "        a.append(A_CSR[1].count(i))\n",
    "        t = t + a[i]\n",
    "        trow.append(tmp)\n",
    "        \n",
    "    for i in np.arange(n):\n",
    "        for j in np.arange(A_CSR[0][i + 1] - A_CSR[0][i]):\n",
    "            tcol[A_CSR[1][A_CSR[0][i] + j]].append(i)\n",
    "            tval[A_CSR[1][A_CSR[0][i] + j]].append(A_CSR[2][A_CSR[0][i] + j])\n",
    "            \n",
    "    return trow, tcol, tval\n",
    "#版本2，数组链表实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 370,
   "id": "9495eef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Transfer(A_CSR):\n",
    "    n = np.size(A_CSR[0]) - 1\n",
    "    m = np.max(A_CSR[1]) + 1\n",
    "    \n",
    "    if n <= 0 or m <= 0:\n",
    "        return abort(2)\n",
    "    \n",
    "    a = []\n",
    "    t = 0\n",
    "    t_row = [0]\n",
    "    t_col = np.ones(np.size(A_CSR[1])) * (-1)\n",
    "    t_val = np.ones(np.size(A_CSR[1])) \n",
    "    \n",
    "    for i in np.arange(m):\n",
    "        a.append(A_CSR[1].count(i))\n",
    "        t = t + a[i]\n",
    "        t_row.append(t)\n",
    "\n",
    "    for i in np.arange(n):\n",
    "        for j in np.arange(A_CSR[0][i + 1] - A_CSR[0][i]):\n",
    "            k = 0\n",
    "            while t_col[t_row[A_CSR[1][A_CSR[0][i] + j]] + k] != -1:\n",
    "                k = k + 1\n",
    "            t_col[t_row[A_CSR[1][A_CSR[0][i] + j]] + k] = i\n",
    "            t_val[t_row[A_CSR[1][A_CSR[0][i] + j]] + k] = A_CSR[2][A_CSR[0][i] + j]\n",
    "            \n",
    "    T_CSR = [t_row, list(map(int, t_col)), list(t_val)]\n",
    "    \n",
    "    return T_CSR\n",
    "#最终版本，数组实现（思路为索引）\n",
    "#注意如果无零行or零列，则n与m无需原矩阵给出，可由按行存储唯一确定\n",
    "#如有某行or列有非零元，则事先需用compress_universal压缩，否则大小与原矩阵不符合，由计算得出的矩阵也是如此！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 371,
   "id": "42b9276f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Multiply(A_CSR, B_CSR):\n",
    "    n = np.size(A_CSR[0]) - 1\n",
    "    m = np.max(A_CSR[1]) + 1\n",
    "    p = np.size(B_CSR[0]) - 1\n",
    "    q = np.max(B_CSR[1]) + 1\n",
    "    \n",
    "    if n <= 0 or m <= 0 or p <= 0 or q <= 0:\n",
    "        return abort(2)\n",
    "    if m != p:\n",
    "        return abort(1)\n",
    "    \n",
    "    c_row = np.zeros(n + 1)\n",
    "    c_col = []\n",
    "    c_val = []\n",
    "    \n",
    "    for i in np.arange(n):\n",
    "        t_val = np.zeros(q)\n",
    "        for j in np.arange((A_CSR[0][i + 1] - A_CSR[0][i])):\n",
    "            t_row = A_CSR[1][A_CSR[0][i] + j] \n",
    "            for k in np.arange(B_CSR[0][t_row + 1] - B_CSR[0][t_row]):\n",
    "                t_col = B_CSR[1][B_CSR[0][t_row] + k]\n",
    "                t_val[t_col] += A_CSR[2][A_CSR[0][i] + j] * B_CSR[2][B_CSR[0][t_row] + k]\n",
    "                \n",
    "        t = 0\n",
    "        for l in np.arange(q):\n",
    "            if(t_val[l] > eps or t_val[l] < -eps):\n",
    "                c_col.append(l)\n",
    "                c_val.append(t_val[l])\n",
    "                t = t + 1\n",
    "        c_row[i + 1] = c_row[i] + t  \n",
    "        \n",
    "    C_CSR = [list(map(int, c_row)), c_col, c_val]\n",
    "    \n",
    "    return C_CSR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 373,
   "id": "fd58dc15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[3., 1., 0., 0., 0.],\n",
       "        [2., 3., 1., 0., 0.],\n",
       "        [0., 2., 3., 1., 0.],\n",
       "        [0., 0., 2., 3., 1.],\n",
       "        [0., 0., 0., 2., 3.]]),\n",
       " array([[4., 0., 0., 0., 0.],\n",
       "        [0., 4., 0., 1., 0.],\n",
       "        [0., 0., 0., 0., 0.],\n",
       "        [0., 3., 0., 4., 0.],\n",
       "        [0., 0., 0., 0., 4.]]),\n",
       " array([[2., 1., 0., 0.],\n",
       "        [1., 2., 1., 0.],\n",
       "        [0., 1., 2., 1.],\n",
       "        [0., 0., 1., 2.]]))"
      ]
     },
     "execution_count": 373,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A, B, C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 374,
   "id": "cd9005ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([0, 2, 5, 8, 11, 13],\n",
       " [0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4],\n",
       " [3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0])"
      ]
     },
     "execution_count": 374,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_CSR = Compress(A)\n",
    "C_CSR = Compress(C)\n",
    "A_CSR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 375,
   "id": "9feafea3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([0, 1, 3, 4, 6, 7],\n",
       " [0, 1, 3, 2, 1, 3, 4],\n",
       " [4.0, 4.0, 1.0, 0.0, 3.0, 4.0, 4.0])"
      ]
     },
     "execution_count": 375,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B_CSR = Compress_universal(B)\n",
    "B_CSR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 376,
   "id": "47800ac2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The size of matrix is wrong!',\n",
       " array([[3., 1., 0., 0., 0., 0., 0.],\n",
       "        [2., 3., 1., 0., 0., 0., 0.],\n",
       "        [0., 2., 3., 1., 0., 0., 0.],\n",
       "        [0., 0., 2., 3., 1., 0., 0.],\n",
       "        [0., 0., 0., 2., 3., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0.]]))"
      ]
     },
     "execution_count": 376,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unCompress(A_CSR, 3, 3), unCompress(A_CSR, 7, 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 385,
   "id": "8bd39979",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([[0, 1, 3, 4, 6, 7],\n",
       "  [0, 1, 3, 2, 1, 3, 4],\n",
       "  [4.0, 4.0, 3.0, 0.0, 1.0, 4.0, 4.0]],\n",
       " array([[4., 0., 0., 0., 0.],\n",
       "        [0., 4., 0., 3., 0.],\n",
       "        [0., 0., 0., 0., 0.],\n",
       "        [0., 1., 0., 4., 0.],\n",
       "        [0., 0., 0., 0., 4.]]))"
      ]
     },
     "execution_count": 385,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Transfer(B_CSR),unCompress(Transfer(B_CSR), 5, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 378,
   "id": "80010623",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The sizes of matrixs are not matched!'"
      ]
     },
     "execution_count": 378,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Multiply(A_CSR, C_CSR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 379,
   "id": "d181249b",
   "metadata": {},
   "outputs": [],
   "source": [
    "D_CSR = Multiply(A_CSR, B_CSR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 380,
   "id": "77a07d3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[12.,  4.,  0.,  1.,  0.],\n",
       "       [ 8., 12.,  0.,  3.,  0.],\n",
       "       [ 0., 11.,  0.,  6.,  0.],\n",
       "       [ 0.,  9.,  0., 12.,  4.],\n",
       "       [ 0.,  6.,  0.,  8., 12.]])"
      ]
     },
     "execution_count": 380,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unCompress(D_CSR, 5, 5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
